# -*- coding: utf-8 -*-
# @Time    : 2021/4/19
# @File    : stylizer.py
import os

from PIL import Image
from models import TransformerNet
from torch.autograd import Variable
from torchvision.utils import save_image

from utils import *

# image_path", type=str, required=True, help="Path to image"
# checkpoint_model", type=str, required=True, help="Path to checkpoint model"

args = {
	'image_path': './images/content/3.jpg',
	'checkpoint_model': './checkpoints/3_4000.pth',
}

if __name__ == "__main__":
	os.makedirs("images/outputs", exist_ok=True)

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	transform = style_transform(None)

	# Define model and load model checkpoint
	transformer = TransformerNet().to(device)
	transformer.load_state_dict(torch.load(args['checkpoint_model']))
	transformer.eval()

	# Prepare input
	image_tensor = Variable(transform(Image.open(args['image_path']))).to(device)
	image_tensor = image_tensor.unsqueeze(0)

	# Stylize image
	with torch.no_grad():
		stylized_image = denormalize(transformer(image_tensor)).cpu()

	# Save image
	fn = args['image_path'].split("/")[-1]
	save_image(stylized_image, f"images/outputs/stylized-{fn}")
